import pygame
import numpy as np
import math
import random
import time
from enum import Enum
from collections import deque


# ===== 1. Core Data Types =====
class SensorType(Enum):
    RADAR = 1
    CAMERA = 2


class ObjectClass(Enum):
    VEHICLE = 1
    NON_VEHICLE = 2
    UNKNOWN = 3


class SensorObject:
    def __init__(self, id, distance, obj_class, confidence,
                 angle=0.0, speed=0.0, pos_x=0.0, pos_y=0.0):
        self.id = id
        self.distance = distance
        self.obj_class = obj_class
        self.confidence = confidence
        self.angle = angle
        self.speed = speed
        self.pos_x = pos_x
        self.pos_y = pos_y

    def to_dict(self):
        return {
            'id': self.id,
            'distance': self.distance,
            'class': self.obj_class,
            'confidence': self.confidence,
            'angle': self.angle,
            'speed': self.speed,
            'pos_x': self.pos_x,
            'pos_y': self.pos_y
        }


class SensorData:
    def __init__(self, sensor_type, timestamp, objects):
        self.sensor_type = sensor_type
        self.timestamp = timestamp
        self.objects = objects


# ===== 2. D-S Fusion Core =====
class DSFusionCore:
    def __init__(self):
        self.frame = [ObjectClass.VEHICLE, ObjectClass.NON_VEHICLE]
        self.logger = RealTimeLogger()

    def compute_bpa(self, sensor_data):
        bpa_list = []
        for obj in sensor_data.objects:
            obj_dict = obj.to_dict()

            if sensor_data.sensor_type == SensorType.RADAR:
                is_dynamic = obj_dict['speed'] > 0.5
                bpa = {
                    ObjectClass.VEHICLE: 0.9 if is_dynamic else 0.7,
                    ObjectClass.NON_VEHICLE: 0.05,
                    'Θ': 0.05 if is_dynamic else 0.25
                }
            else:  # CAMERA
                if obj_dict['class'] == ObjectClass.VEHICLE:
                    bpa = {
                        ObjectClass.VEHICLE: min(0.95, obj_dict['confidence']),
                        ObjectClass.NON_VEHICLE: 0.05,
                        'Θ': 0.05
                    }
                else:
                    if random.random() < 0.1:
                        bpa = {
                            ObjectClass.VEHICLE: 0.7,
                            ObjectClass.NON_VEHICLE: 0.2,
                            'Θ': 0.1
                        }
                    else:
                        bpa = {
                            ObjectClass.VEHICLE: 0.05,
                            ObjectClass.NON_VEHICLE: min(0.9, obj_dict['confidence']),
                            'Θ': 0.05
                        }
            bpa_list.append((obj_dict['id'], bpa))
        return bpa_list

    def dempster_combine(self, bpa1, bpa2, obj_id):
        K = 0
        for A in bpa1:
            for B in bpa2:
                if A != B and A != 'Θ' and B != 'Θ':
                    K += bpa1[A] * bpa2[B]

        if K > 0.5:
            self.logger.add(f"Warning: Target {obj_id} evidence conflict (K={K:.2f})", "WARN")

        combined_bpa = {}
        for A in bpa1:
            for B in bpa2:
                if set([A, B]) in [{'Θ'}, {ObjectClass.VEHICLE}, {ObjectClass.NON_VEHICLE}]:
                    key = A if A != 'Θ' else B
                    value = bpa1[A] * bpa2[B]
                    combined_bpa[key] = combined_bpa.get(key, 0) + value

        for key in combined_bpa:
            combined_bpa[key] /= (1 - K)
        return combined_bpa


# ===== 3. Enhanced Data Generator =====
class RealisticDataGenerator:
    def __init__(self):
        self.obj_counter = 1
        self.last_positions = {
            1: [0, 78],  # Moving vehicle
            2: [0, 118],  # Static vehicle
            3: [0, 100]  # Metal sign
        }

    def _add_sensor_error(self, true_value, sensor_type, is_angle=False):
        if sensor_type == SensorType.RADAR:
            if is_angle:
                return true_value + random.gauss(0, 0.5)
            else:
                return true_value + random.uniform(0.2, 0.5)
        else:  # CAMERA
            distance_factor = min(1, true_value / 200)
            if is_angle:
                return true_value + random.gauss(0, 1) * distance_factor
            else:
                return true_value + random.gauss(0, 0.3) * distance_factor

    def generate_radar_data(self, timestamp):
        objects = []
        self.last_positions[1][0] += 2.0

        # Moving vehicle (ID=1)
        true_dist = math.sqrt(self.last_positions[1][0] ** 2 + self.last_positions[1][1] ** 2)
        true_angle = math.degrees(math.atan2(self.last_positions[1][0], self.last_positions[1][1]))
        objects.append(SensorObject(
            id=1,
            distance=self._add_sensor_error(true_dist, SensorType.RADAR),
            obj_class=ObjectClass.VEHICLE,
            confidence=0.92,
            angle=self._add_sensor_error(true_angle, SensorType.RADAR, True),
            speed=random.gauss(10, 0.3),
            pos_x=self.last_positions[1][0],
            pos_y=self.last_positions[1][1]
        ))

        # Static vehicle (ID=2)
        objects.append(SensorObject(
            id=2,
            distance=self._add_sensor_error(118.7, SensorType.RADAR),
            obj_class=ObjectClass.VEHICLE,
            confidence=0.85,
            angle=random.gauss(-0.5, 0.3),
            speed=0,
            pos_x=0,
            pos_y=118.7
        ))

        # Metal sign (ID=3)
        objects.append(SensorObject(
            id=3,
            distance=self._add_sensor_error(100.5, SensorType.RADAR),
            obj_class=ObjectClass.NON_VEHICLE,
            confidence=0.78,
            angle=random.gauss(2.0, 0.5),
            speed=0,
            pos_x=3.5,
            pos_y=100.2
        ))

        return SensorData(SensorType.RADAR, timestamp, objects)

    def generate_camera_data(self, timestamp):
        objects = []
        # Moving vehicle
        true_dist = math.sqrt(self.last_positions[1][0] ** 2 + self.last_positions[1][1] ** 2)
        true_angle = math.degrees(math.atan2(self.last_positions[1][0], self.last_positions[1][1]))
        objects.append(SensorObject(
            id=1,
            distance=self._add_sensor_error(true_dist, SensorType.CAMERA),
            obj_class=ObjectClass.VEHICLE,
            confidence=0.92,
            angle=self._add_sensor_error(true_angle, SensorType.CAMERA, True),
            speed=0,
            pos_x=self.last_positions[1][0],
            pos_y=self.last_positions[1][1]
        ))

        # Static vehicle
        objects.append(SensorObject(
            id=2,
            distance=self._add_sensor_error(118.7, SensorType.CAMERA),
            obj_class=ObjectClass.VEHICLE,
            confidence=0.85,
            angle=random.gauss(-0.5, 0.3),
            speed=0,
            pos_x=0,
            pos_y=118.7
        ))

        # Metal sign (10% misclassification)
        obj_class = ObjectClass.VEHICLE if random.random() < 0.1 else ObjectClass.NON_VEHICLE
        confidence = random.uniform(0.6, 0.8) if obj_class == ObjectClass.VEHICLE else 0.95
        objects.append(SensorObject(
            id=3,
            distance=self._add_sensor_error(99.2, SensorType.CAMERA),
            obj_class=obj_class,
            confidence=confidence,
            angle=random.gauss(1.8, 0.8),
            speed=0,
            pos_x=3.5,
            pos_y=100.2
        ))

        return SensorData(SensorType.CAMERA, timestamp, objects)


# ===== 4. Enhanced Visualization with BEV =====
class TrailVisualization:
    def __init__(self, width=1200, height=800):
        pygame.init()
        self.screen = pygame.display.set_mode((width, height))
        pygame.display.set_caption("Multi-Sensor Fusion Tracking System")
        self.font = pygame.font.SysFont('Arial', 16)
        self.width, self.height = width, height
        self.bev_scale_main = 5  # Main view scale
        self.ego_x, self.ego_y = width // 2, height - 50
        self.trail_length = 30

        # BEV parameters
        self.camera_mode = "FPS"  # FPS or BEV
        self.bev_surface = pygame.Surface((width // 2, height // 2))  # Separate BEV layer
        self.bev_scale = 8  # Meters to pixels for BEV
        self.bev_margin = 20  # Margin for BEV
        self.min_x, self.max_x = -50, 150  # World coordinate range
        self.min_y, self.max_y = -50, 200

        # Trail history
        self.radar_trails = {1: deque(maxlen=self.trail_length),
                             2: deque(maxlen=self.trail_length),
                             3: deque(maxlen=self.trail_length)}
        self.camera_trails = {1: deque(maxlen=self.trail_length),
                              2: deque(maxlen=self.trail_length),
                              3: deque(maxlen=self.trail_length)}
        self.fusion_trails = {1: deque(maxlen=self.trail_length),
                              2: deque(maxlen=self.trail_length),
                              3: deque(maxlen=self.trail_length)}

    def world_to_bev(self, x, y):
        """Convert world coordinates to BEV coordinates (scaling + centering)"""
        bev_x = int((x - self.min_x) * self.bev_scale) + self.bev_margin
        bev_y = int((self.max_y - y) * self.bev_scale) + self.bev_margin
        return bev_x, bev_y

    def _calc_position(self, obj_dict):
        try:
            angle_rad = math.radians(obj_dict['angle'])
            x = self.ego_x + obj_dict['distance'] * self.bev_scale_main * math.sin(angle_rad)
            y = self.ego_y - obj_dict['distance'] * self.bev_scale_main * math.cos(angle_rad)
            return x, y
        except (TypeError, KeyError) as e:
            print(f"Position calculation error: {e}")
            return self.ego_x, self.ego_y

    def draw_trail(self, trail, color, width=2):
        if len(trail) > 1:
            points = [(int(x), int(y)) for x, y in trail]
            pygame.draw.lines(self.screen, color, False, points, width)

    def draw_object(self, obj, color, label):
        obj_dict = obj.to_dict()
        x, y = self._calc_position(obj_dict)

        pygame.draw.circle(self.screen, color, (int(x), int(y)), 8)

        text = self.font.render(f"{label} ID:{obj_dict['id']}", True, color)
        self.screen.blit(text, (int(x) + 12, int(y) - 10))

    def draw_ds_panel(self, results):
        y_offset = 30
        for obj_id, bpa in results.items():
            vehicle_conf = bpa.get(ObjectClass.VEHICLE, 0)
            non_vehicle_conf = bpa.get(ObjectClass.NON_VEHICLE, 0)

            color = (255, 255, 0) if vehicle_conf > 0.8 else (200, 200, 200)
            text = self.font.render(
                f"Target{obj_id}: Vehicle={vehicle_conf:.2f} Non-Vehicle={non_vehicle_conf:.2f}",
                True, color
            )
            self.screen.blit(text, (20, y_offset))
            y_offset += 25

    def render_first_person_view(self, radar_data):
        """First-person perspective rendering (simplified)"""
        self.screen.fill((0, 10, 20))  # Dark background

        # Draw road vanishing point
        pygame.draw.line(self.screen, (100, 100, 100),
                         (self.ego_x, self.height),
                         (self.ego_x, self.ego_y - 200), 3)

        # Draw radar targets
        for obj in radar_data.objects:
            obj_dict = obj.to_dict()
            # Perspective projection: x_screen = (obj_x / distance) * focal_length
            focal_length = 300
            screen_x = self.ego_x + (obj_dict['pos_x'] / obj_dict['distance']) * focal_length
            screen_y = self.ego_y - 50  # Fixed horizon position

            color = (255, 0, 0) if obj_dict['class'] == ObjectClass.VEHICLE else (0, 255, 0)
            pygame.draw.rect(self.screen, color, (screen_x - 5, screen_y - 5, 10, 10))

            # Draw distance label
            dist_text = self.font.render(f"{obj_dict['distance']:.1f}m", True, (200, 200, 200))
            self.screen.blit(dist_text, (screen_x - 15, screen_y - 25))

    def render_bev(self):
        """Render Bird's Eye View (Core BEV logic)"""
        self.bev_surface.fill((30, 30, 50))  # Dark blue background

        # Draw road grid
        for i in range(0, 200, 10):
            bev_x1, bev_y1 = self.world_to_bev(self.min_x, i)
            bev_x2, bev_y2 = self.world_to_bev(self.max_x, i)
            pygame.draw.line(self.bev_surface, (60, 60, 80), (bev_x1, bev_y1), (bev_x2, bev_y2), 1)

        # Draw ego vehicle (BEV center)
        ego_bev_x, ego_bev_y = self.world_to_bev(0, 0)
        pygame.draw.rect(self.bev_surface, (0, 150, 255),
                         (ego_bev_x - 5, ego_bev_y - 5, 10, 10))

        # Draw object trails
        for obj_id in [1, 2, 3]:
            # Radar trails (red)
            if len(self.radar_trails[obj_id]) > 1:
                points = [self.world_to_bev(x, y) for x, y in self.radar_trails[obj_id]]
                pygame.draw.lines(self.bev_surface, (200, 50, 50), False, points, 2)

            # Camera trails (green)
            if len(self.camera_trails[obj_id]) > 1:
                points = [self.world_to_bev(x, y) for x, y in self.camera_trails[obj_id]]
                pygame.draw.lines(self.bev_surface, (50, 200, 50), False, points, 2)

            # Fusion trails (yellow)
            if len(self.fusion_trails[obj_id]) > 1:
                points = [self.world_to_bev(x, y) for x, y in self.fusion_trails[obj_id]]
                pygame.draw.lines(self.bev_surface, (255, 255, 0), False, points, 3)

        # Draw current positions
        for obj_id, trail in self.fusion_trails.items():
            if trail:
                x, y = trail[-1]
                bev_x, bev_y = self.world_to_bev(x, y)
                pygame.draw.circle(self.bev_surface, (255, 255, 0), (bev_x, bev_y), 5)

        # Draw coordinate axis
        pygame.draw.line(self.bev_surface, (200, 200, 200),
                         self.world_to_bev(0, 0),
                         self.world_to_bev(50, 0), 3)  # X-axis
        pygame.draw.line(self.bev_surface, (200, 200, 200),
                         self.world_to_bev(0, 0),
                         self.world_to_bev(0, 50), 3)  # Y-axis

        # Add BEV title
        title = self.font.render("Bird's Eye View (BEV)", True, (255, 255, 255))
        self.bev_surface.blit(title, (10, 5))

        # Overlay to main screen
        self.screen.blit(self.bev_surface, (self.width - self.bev_surface.get_width() - 10, 10))

    def render(self, radar_data, camera_data, fusion_results, logger):
        # Clear screen
        self.screen.fill((0, 20, 40))

        # === Render main view ===
        if self.camera_mode == "FPS":
            self.render_first_person_view(radar_data)
        else:
            # Render BEV as main view
            self.render_bev()

        # === Draw BEV minimap (always visible) ===
        minimap_width, minimap_height = 300, 200
        minimap_surface = pygame.Surface((minimap_width, minimap_height))
        minimap_surface.fill((30, 30, 50))

        # Draw ego vehicle at center
        pygame.draw.rect(minimap_surface, (0, 150, 255),
                         (minimap_width // 2 - 3, minimap_height // 2 - 3, 6, 6))

        # Draw objects (simplified)
        for obj in radar_data.objects:
            obj_dict = obj.to_dict()
            dx = int(obj_dict['pos_x'] * 1.5)  # Simplified scaling
            dy = int(obj_dict['pos_y'] * 1.5)
            color = (255, 0, 0) if obj_dict['class'] == ObjectClass.VEHICLE else (0, 255, 0)
            pygame.draw.circle(minimap_surface, color,
                               (minimap_width // 2 + dx, minimap_height // 2 - dy), 3)

        self.screen.blit(minimap_surface, (self.width - minimap_width - 20, self.height - minimap_height - 20))

        # === Draw information panels ===
        self.draw_ds_panel(fusion_results)
        logger.draw(self.screen, self.font, 20, 500)

        # Draw mode indicator
        mode_text = self.font.render(f"View Mode: {self.camera_mode} (Press V to toggle)", True, (255, 255, 0))
        self.screen.blit(mode_text, (self.width - 350, 50))

        pygame.display.flip()

    def handle_events(self):
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_v:  # Toggle view mode
                    self.camera_mode = "BEV" if self.camera_mode == "FPS" else "FPS"
                elif event.key == pygame.K_ESCAPE:
                    return False
        return True


# ===== 5. Logging System =====
class RealTimeLogger:
    def __init__(self, max_lines=8):
        self.messages = []
        self.max_lines = max_lines
        self.colors = {"INFO": (200, 200, 200), "WARN": (255, 200, 0), "ERROR": (255, 50, 50)}

    def add(self, msg, level="INFO"):
        timestamp = time.strftime("%H:%M:%S")
        self.messages.append((f"[{timestamp}] {msg}", level))
        self.messages = self.messages[-self.max_lines:]

    def draw(self, screen, font, x, y):
        pygame.draw.rect(screen, (0, 0, 0, 180), (x - 10, y - 10, 600, 25 * self.max_lines + 10))
        for i, (msg, level) in enumerate(self.messages):
            text = font.render(msg, True, self.colors.get(level, (200, 200, 200)))
            screen.blit(text, (x, y + i * 25))


# ===== 6. Main System =====
def main():
    vis = TrailVisualization()
    data_gen = RealisticDataGenerator()
    ds_fusion = DSFusionCore()
    logger = RealTimeLogger()

    logger.add("System started: Multi-sensor fusion initialized", "INFO")
    logger.add("Press V to toggle Bird's Eye View", "INFO")

    timestamp = 0
    frame_count = 0

    running = True
    clock = pygame.time.Clock()

    while running:
        timestamp += 100
        frame_count += 1

        # Handle events
        running = vis.handle_events()

        # Generate data
        radar_data = data_gen.generate_radar_data(timestamp)
        camera_data = data_gen.generate_camera_data(timestamp)
        logger.add(
            f"Frame{frame_count}: Radar targets: {len(radar_data.objects)} | Camera targets: {len(camera_data.objects)}")

        # Fusion processing
        fusion_results = {}
        radar_bpa = ds_fusion.compute_bpa(radar_data)
        camera_bpa = ds_fusion.compute_bpa(camera_data)

        for (radar_id, r_bpa), (cam_id, c_bpa) in zip(radar_bpa, camera_bpa):
            fused_bpa = ds_fusion.dempster_combine(r_bpa, c_bpa, radar_id)
            fusion_results[radar_id] = fused_bpa
            if fused_bpa.get(ObjectClass.VEHICLE, 0) > 0.85:
                logger.add(f"Target{radar_id}: Vehicle confidence {100 * fused_bpa[ObjectClass.VEHICLE]:.0f}%", "INFO")

        # Update trails
        for obj in radar_data.objects:
            obj_dict = obj.to_dict()
            if obj_dict['id'] in [1, 2, 3]:
                vis.radar_trails[obj_dict['id']].append((obj_dict['pos_x'], obj_dict['pos_y']))

        for obj in camera_data.objects:
            obj_dict = obj.to_dict()
            if obj_dict['id'] in [1, 2, 3]:
                vis.camera_trails[obj_dict['id']].append((obj_dict['pos_x'], obj_dict['pos_y']))

        # Update fusion trails
        for obj_id, bpa in fusion_results.items():
            if bpa.get(ObjectClass.VEHICLE, 0) > 0.7:
                radar_obj = next((o for o in radar_data.objects if o.id == obj_id), None)
                if radar_obj:
                    obj_dict = radar_obj.to_dict()
                    vis.fusion_trails[obj_id].append((obj_dict['pos_x'], obj_dict['pos_y']))

        # Render
        vis.render(radar_data, camera_data, fusion_results, logger)
        clock.tick(10)

    pygame.quit()


if __name__ == "__main__":
    main()